/**@file
   Entry module for kernel stub module.
 
   @par license
   Copyright 2012 Lu, Ken. All rights reserved.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
 
   @author: Lu, Ken (bluewish.ken.lu@gmail.com)
 **/

#include "internal.h"
#include "kstub_if.h"

#ifdef ALLOC_PRAGMA
    #pragma alloc_text(INIT, DriverEntry)
    #pragma alloc_text(PAGE, KSDriverUnload)
    #pragma alloc_text(PAGE, KSEvtDriverContextCleanup)
    #pragma alloc_text(PAGE, KSEvtDeviceShutdown)
    #pragma alloc_text(PAGE, KSEvtDeviceCreate)
    #pragma alloc_text(PAGE, KSEvtDeviceClose)
    #pragma alloc_text(PAGE, KSEvtIoDeviceControl)
    #pragma alloc_text(PAGE, KSEvtIoRead)
    #pragma alloc_text(PAGE, KSEvtIoWrite)
    #pragma alloc_text(PAGE, KSDeviceAdd)
#endif

//
// Function Implementation
//
/**
 * This routine gets called by the system to initialize the driver.
 * 
 * @author klu2 (2012/3/31)
 * 
 * @param pDriverObject     a pointer to the object that represents 
 *                          this device driver.
 * @param pRegistryPath     a pointer to our Services key in the 
 *                          registry.
 * 
 * @return NTSTATUS 
 */
NTSTATUS
    DriverEntry(
    PDRIVER_OBJECT  pDriverObject,
    PUNICODE_STRING pRegistryPath
    )
{
    WDF_DRIVER_CONFIG       config;
    WDF_OBJECT_ATTRIBUTES   attributes;
    WDFDRIVER               hDriver;
    NTSTATUS                status;
    PWDFDEVICE_INIT         pInit;

    UNREFERENCED_PARAMETER(pDriverObject);
    UNREFERENCED_PARAMETER(pRegistryPath);

    k_func_enter();
    k_debug("Build: %s %s\n", __DATE__, __TIME__);

    WDF_DRIVER_CONFIG_INIT(&config, WDF_NO_EVENT_CALLBACK );

    //
    // Tell the framework that this is non-pnp driver so that it doesn't
    // set the default AddDevice routine.
    //
    config.DriverInitFlags |= WdfDriverInitNonPnpDriver;

    //
    // NonPnp driver must explicitly register an unload routine for
    // the driver to be unloaded.
    //
    config.EvtDriverUnload = KSDriverUnload;

    WDF_OBJECT_ATTRIBUTES_INIT(&attributes);
    attributes.EvtCleanupCallback = KSEvtDriverContextCleanup;

    status = WdfDriverCreate(pDriverObject, 
                 pRegistryPath, 
                 &attributes, 
                 &config, 
                 &hDriver);
    if (!NT_SUCCESS(status)) {
        k_debug("Fail to create wdf device: status=0x%x\n", status);
        return status;
    }

    //
    // In order to create a control device, we first need to allocate a
    // WDFDEVICE_INIT structure and set all properties.
    //
    pInit = WdfControlDeviceInitAllocate(hDriver, &SDDL_DEVOBJ_SYS_ALL_ADM_RWX_WORLD_RW_RES_R);
    if (pInit == NULL) {
        k_debug("Fail to allocate wdf init structure.\n");
        return STATUS_INSUFFICIENT_RESOURCES;
    }


    status = KSDeviceAdd(hDriver, pInit);
    k_func_leave();

    return status;
}

VOID 
    KSDriverUnload(
    WDFDRIVER hDriver
    )
{
    UNREFERENCED_PARAMETER(hDriver);
    PAGED_CODE();
    k_func_enter();

    k_func_leave();
}

VOID
    KSEvtDriverContextCleanup(
    IN WDFDRIVER hDriver
    )
{
    UNREFERENCED_PARAMETER(hDriver);
    PAGED_CODE();
    k_func_enter();

    //
    // No need to free the controldevice object explicitly because it will
    // be deleted when the Driver object is deleted due to the default parent
    // child relationship between Driver and ControlDevice.
    //

    k_func_leave();
}

NTSTATUS
    KSDeviceAdd(
    WDFDRIVER           hDriver,
    PWDFDEVICE_INIT     pInit
    )
{
    NTSTATUS                status;
    WDF_FILEOBJECT_CONFIG   fileConfig;
    WDF_OBJECT_ATTRIBUTES   attributes;
    WDFDEVICE               hDevice;
    WDF_IO_QUEUE_CONFIG     ioQueueConfig;
    WDFQUEUE                queue;
    DECLARE_CONST_UNICODE_STRING(ntDeviceName, KS_NT_DEVICE_NAME);
    DECLARE_CONST_UNICODE_STRING(dosDeviceName, KS_DOS_DEVICE_NAME);

    UNREFERENCED_PARAMETER(hDriver);

    PAGED_CODE();

    WdfDeviceInitSetExclusive(pInit, TRUE);

    WdfDeviceInitSetIoType(pInit, WdfDeviceIoBuffered);

    status = WdfDeviceInitAssignName(pInit, &ntDeviceName);

    if (!NT_SUCCESS(status)) {
        goto ret;
    }

    WdfControlDeviceInitSetShutdownNotification(
        pInit,
        KSEvtDeviceShutdown,
        WdfDeviceShutdown);

    //
    // Initialize WDF_FILEOBJECT_CONFIG_INIT struct to tell the
    // framework whether you are interested in handling Create, Close and
    // Cleanup requests that gets generated when an application or another
    // kernel component opens an handle to the device. If you don't register
    // the framework default behaviour would be to complete these requests
    // with STATUS_SUCCESS. A driver might be interested in registering these
    // events if it wants to do security validation and also wants to maintain
    // per handle (fileobject) context.
    //
    WDF_FILEOBJECT_CONFIG_INIT(
        &fileConfig,
        KSEvtDeviceCreate,
        KSEvtDeviceClose,
        WDF_NO_EVENT_CALLBACK // not interested in Cleanup
        );

    WdfDeviceInitSetFileObjectConfig(
        pInit,
        &fileConfig,
        WDF_NO_OBJECT_ATTRIBUTES
        );

    WDF_OBJECT_ATTRIBUTES_INIT_CONTEXT_TYPE(&attributes, KS_DEVICE_EXTENSION);

    status = WdfDeviceCreate(&pInit, &attributes, &hDevice);
    if (!NT_SUCCESS(status)) {
        k_debug("Fail to create device object. status=0x%x\n", status);
        goto ret;
    }


    status = WdfDeviceCreateSymbolicLink(hDevice, &dosDeviceName);
    if (!NT_SUCCESS(status)) {
        k_debug("Fail to create symbol. status=0x%x\n", status);
        goto ret;
    }

    //
    // Configure a default queue so that requests that are not
    // configure-fowarded using WdfDeviceConfigureRequestDispatching to goto
    // other queues get dispatched here.
    //
    WDF_IO_QUEUE_CONFIG_INIT_DEFAULT_QUEUE(
        &ioQueueConfig,
        WdfIoQueueDispatchSequential
        );

    ioQueueConfig.EvtIoRead             = KSEvtIoRead;
    ioQueueConfig.EvtIoWrite            = KSEvtIoWrite;
    ioQueueConfig.EvtIoDeviceControl    = KSEvtIoDeviceControl;

    WDF_OBJECT_ATTRIBUTES_INIT(&attributes);

    //
    // Since we are using Zw function set execution level to passive so that
    // framework ensures that our Io callbacks called at only passive-level
    // even if the request came in at DISPATCH_LEVEL from another driver.
    //
    //attributes.ExecutionLevel = WdfExecutionLevelPassive;

    status = WdfIoQueueCreate(hDevice,
                 &ioQueueConfig,
                 &attributes,
                 &queue // pointer to default queue
                 );
    if (!NT_SUCCESS(status)) {
        k_debug("Fail to create I/O queue. status=0x%x\n", status);
        goto ret;
    }

    //
    // Control devices must notify WDF when they are done initializing.   I/O is
    // rejected until this call is made.
    //
    WdfControlFinishInitializing(hDevice);

    ret:
    if (pInit != NULL) {
        WdfDeviceInitFree(pInit);
    }
    return status;
}

VOID
    KSEvtDeviceShutdown(
    WDFDEVICE  hDevice
    )
{
    UNREFERENCED_PARAMETER(hDevice);
    PAGED_CODE()
    k_func_enter();
    k_func_leave();
    return;
}

VOID
    KSEvtDeviceCreate (
    IN WDFDEVICE            hDevice,
    IN WDFREQUEST           hRequest,
    IN WDFFILEOBJECT        hFileObject
    )
{
    NTSTATUS    status;

    UNREFERENCED_PARAMETER(hDevice);
    UNREFERENCED_PARAMETER(hFileObject);

    PAGED_CODE();
    k_func_enter();

    status = STATUS_SUCCESS;
    WdfRequestComplete(hRequest, status);

    k_func_leave();
}

VOID
    KSEvtDeviceClose (
    IN WDFFILEOBJECT    hFileObject
    )
{
    UNREFERENCED_PARAMETER(hFileObject);
    PAGED_CODE();
    k_func_enter();
    k_func_leave();
}

VOID
    KSEvtIoRead(
    IN WDFQUEUE         hQueue,
    IN WDFREQUEST       hRequest,
    IN size_t           length
    )
{
    UNREFERENCED_PARAMETER(hQueue);
    UNREFERENCED_PARAMETER(length);
    PAGED_CODE();
    k_func_enter();

    WdfRequestCompleteWithInformation(
        hRequest, 
        STATUS_SUCCESS, 
        0);

    k_func_leave();
}

VOID
    KSEvtIoWrite(
    IN WDFQUEUE         hQueue,
    IN WDFREQUEST       hRequest,
    IN size_t           length
    )
{
    UNREFERENCED_PARAMETER(hQueue);
    UNREFERENCED_PARAMETER(hRequest);
    UNREFERENCED_PARAMETER(length);

    PAGED_CODE();
    k_func_enter();

    WdfRequestCompleteWithInformation(
        hRequest, 
        STATUS_SUCCESS, 
        0);

    k_func_leave();

}

VOID
    KSEvtIoDeviceControl(
    IN WDFQUEUE         hQueue,
    IN WDFREQUEST       hRequest,
    IN size_t           outputBufferLength,
    IN size_t           inputBufferLength,
    IN ULONG            ioControlCode
    )
{
    MMIO_READ_PARAM *pMmIoReadParam;
    MMIO_SCAN_PARAM *pMmIoScanParam;
    IO_PORT_PARAM   *pIoPortParam;
    IO_PARAM        *pIoParam;

    PCHAR inbuf = NULL, outbuf = NULL;
    size_t bufsize;
    NTSTATUS status = STATUS_SUCCESS;
    u32_t actualsize = 0;

    UNREFERENCED_PARAMETER(hQueue);
    UNREFERENCED_PARAMETER(outputBufferLength);
    UNREFERENCED_PARAMETER(inputBufferLength);
    UNREFERENCED_PARAMETER(ioControlCode);

    PAGED_CODE();
    k_func_enter();

    k_debug("IO control code: 0x%x, input len: 0x%x, output len: 0x%x", 
        ioControlCode,
        inputBufferLength,
        outputBufferLength);

    if (inputBufferLength != 0) {
        status = WdfRequestRetrieveInputBuffer(
                     hRequest, 
                     0, 
                     &inbuf, 
                     &bufsize);
        if (!NT_SUCCESS(status)) {
            status = STATUS_INSUFFICIENT_RESOURCES;
            goto exit;
        }
        k_assert(bufsize == inputBufferLength);
    }
    if (outputBufferLength != 0) {
        status = WdfRequestRetrieveOutputBuffer(
                     hRequest, 
                     0, 
                     &outbuf, 
                     &bufsize);
        if (!NT_SUCCESS(status)) {
            status = STATUS_INSUFFICIENT_RESOURCES;
            goto exit;
        }
        k_assert(bufsize == outputBufferLength);
    }

    switch (ioControlCode) {
    case KS_DRIVER_MMIO_READ:
        k_assert(inputBufferLength  != 0);
        k_assert(outputBufferLength != 0);
        pMmIoReadParam = (MMIO_READ_PARAM*) inbuf;
        actualsize = outputBufferLength;
        status = k_mmio_read(
                     (un_t)pMmIoReadParam->base, 
                     pMmIoReadParam->length, 
                     (u8_t*)outbuf, 
                     &actualsize
                     );
        WdfRequestSetInformation(
            hRequest, 
            outputBufferLength<actualsize?outputBufferLength:actualsize
            );
        break;

    case KS_DRIVER_IO_READ:
        k_assert(inputBufferLength  != 0);
        k_assert(outputBufferLength != 0);
        pIoParam = (IO_PARAM*) inbuf;
        actualsize = outputBufferLength;
        status = KSIoRead(
                     pIoParam->port, 
                     pIoParam->width, 
                     (u8_t*)outbuf, 
                     &actualsize);
        WdfRequestSetInformation(
            hRequest, 
            outputBufferLength<actualsize?outputBufferLength:actualsize
            );
        break;

    case KS_DRIVER_IO_WRITE:
        k_assert(inputBufferLength  != 0);
        pIoParam = (IO_PARAM*) inbuf;
        status = KSIoWrite(
                     pIoParam->port, 
                     pIoParam->width, 
                     pIoParam->value
                     );
        break;

    case KS_DRIVER_IO_PORT_READ:
        k_assert(inputBufferLength  != 0);
        k_assert(outputBufferLength != 0);
        pIoPortParam = (IO_PORT_PARAM*) inbuf;
        status = KSIoPortRead(
                     pIoPortParam->index_port, 
                     pIoPortParam->data_port, 
                     (u8_t*)outbuf);
        if (NT_SUCCESS(status)) {
            WdfRequestSetInformation(hRequest, 0xFF);
        } else {
            WdfRequestSetInformation(hRequest, 0);
        }
        break;

    case KS_DRIVER_MMIO_SCAN:
        k_assert(inputBufferLength  != 0);
        k_assert(outputBufferLength != 0);
        pMmIoScanParam = (MMIO_SCAN_PARAM*) inbuf;
        status = k_mmio_scan(
                   (un_t)pMmIoScanParam->start, 
                   pMmIoScanParam->length, 
                   pMmIoScanParam->scan_buffer, 
                   pMmIoScanParam->scan_buffer_len,
                   (u32_t*)outbuf);
        if (NT_SUCCESS(status)) {
            WdfRequestSetInformation(hRequest, 4);
        } else {
            WdfRequestSetInformation(hRequest, 0);
        }
        break;
    default:
        k_debug("Invalid IO control code.");
        break;
    }

    exit:
    WdfRequestComplete(hRequest, status);
    k_func_leave();
}
